# Copyright 2020 Makani Technologies LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functions to write controller gains into output files."""

import json
import pprint

from makani.control import system_types
from makani.lib.python import c_helpers

_WING_SERIAL_HELPER = c_helpers.EnumHelper('WingSerial', system_types)


def WriteControllers(generating_file, filename, controllers):
  """Writes gain matrices to configuration file.

  Args:
    generating_file: File calling this function.
    filename: Full path to the gain matrix file to write.
    controllers: OrderedDict of airspeed values and gain matrices.
  """

  docstring = ('  """Returns control gain matrices'
               'for any kite serial number."""')

  with open(filename, 'w') as f:
    f.write('"""Automatically generated controllers.'
            '\n'
            '\nThis file was generated by: '
            + generating_file + '.'
            '\n"""'
            '\n'
            '\nfrom makani.control import control_types as m'
            '\n'
            '\n'
            '\ndef GetControllers(wing_serial):'
            '\n' + docstring)
    for i, c in enumerate(controllers):
      if i == 0:
        f.write('\n  if wing_serial == m.%s:' % c['wing_serial'])
      else:
        f.write('\n'
                '\n  elif wing_serial == m.%s:' % c['wing_serial'])
      for k, v in c.iteritems():
        pyexcept = ''
        if k == 'B_flaps_to_pqr_min_airspeed':
          pyexcept = '  # pylint: disable=invalid-name'
        if k != 'wing_serial':
          f.write('\n    %s = (%s' % (k, pyexcept))
          f.write('\n        ')
          f.write('\n        '.join(pprint.pformat(v, width=74).splitlines()))
          f.write('\n    )')

    f.write('\n'
            '\n  else:')
    f.write('\n    assert False, ')
    f.write("'wing_serial %d was not recognized' % wing_serial")

    f.write('\n'
            '\n  return {')
    for k in controllers[0].iterkeys():
      if k != 'wing_serial':
        f.write("\n      '%s': ("
                "\n          %s)," % (k, k))
    f.write('\n  }\n')


def WriteControllersToJson(filename, controllers):
  """Writes gain matrices to json file.

  Args:
    filename: Full path to the json file to write.
    controllers: OrderedDict of airspeed values and gain matrices.
  """
  with open(filename, 'w') as f:
    output_string = json.dumps(controllers, indent=2)
    f.write(output_string)

